#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
问题加载测试

测试各个检索器的问题加载功能
"""

import sys
import os

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 在导入任何库之前强制设置环境变量
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
os.environ["HUGGINGFACE_CO_URL_HOME"] = "https://hf-mirror.com"
os.environ["HUGGINGFACE_HUB_CACHE"] = os.path.expanduser("~/.cache/huggingface")
os.environ["SENTENCE_TRANSFORMERS_HOME"] = os.path.expanduser("~/.cache/sentence_transformers")

# 强制导入配置，确保镜像环境变量已设置
import questionretrieval.config
from loguru import logger

logger.info("🌐 使用 HuggingFace 国内镜像配置")
logger.info(f"📁 缓存目录: {questionretrieval.config.HF_HUB_CACHE}")
logger.info(f"🔗 镜像地址: {questionretrieval.config.HF_MIRROR_ENDPOINT}")

import pytest
from questionretrieval import load_questions

# 测试数据
TEST_QUESTIONS = [
    "大伙房水库的主管部门是哪个单位？",
    "大伙房水库总库容/防洪库容/兴利库容/死库容？",
    "大伙房水库调度目标和任务是什么？",
    "大伙房水库防洪调度的核心任务是什么？",
    "大伙房水库供水调度的核心任务是什么？"
]

def get_test_questions():
    """
    获取测试问题
    """
    try:
        questions = load_questions()
        if len(questions) > 0:
            return questions
    except:
        pass
    return TEST_QUESTIONS

def test_sentence_transformer_load_questions():
    """
    测试 SentenceTransformer 检索器加载问题
    """
    logger.info("\n=== 测试 SentenceTransformer 加载问题 ===")
    
    try:
        from questionretrieval import SentenceTransformerRetriever
        retriever = SentenceTransformerRetriever()
        questions = get_test_questions()
        
        retriever.load_candidates(questions)
        return True
    except ImportError:
        logger.warning("✗ SentenceTransformer 不可用")
        return False
    except Exception as e:
        logger.error(f"✗ 加载问题失败: {e}")
        return False

def test_bert_chinese_load_questions():
    """
    测试 BERT 中文检索器加载问题
    """
    logger.info("\n=== 测试 BERT 中文检索器加载问题 ===")
    
    try:
        from questionretrieval import BertChineseRetriever
        retriever = BertChineseRetriever()
        questions = get_test_questions()
        
        retriever.load_candidates(questions)
        return True
    except ImportError:
        logger.warning("✗ BertChineseRetriever 不可用")
        return False
    except Exception as e:
        logger.error(f"✗ 加载问题失败: {e}")
        return False

def test_bert_similarity_load_questions():
    """
    测试 BERT 相似度检索器加载问题
    """
    logger.info("\n=== 测试 BERT 相似度检索器加载问题 ===")
    
    try:
        from questionretrieval import BertSimilarityRetriever
        retriever = BertSimilarityRetriever()
        questions = get_test_questions()
        
        retriever.load_candidates(questions)
        return True
    except ImportError:
        logger.warning("✗ BertSimilarityRetriever 不可用")
        return False
    except Exception as e:
        logger.error(f"✗ 加载问题失败: {e}")
        return False

def test_all_retrievers_load_questions():
    """
    测试所有检索器加载问题
    """
    logger.info("\n" + "="*50)
    logger.info("检索器问题加载测试")
    logger.info("="*50)
    
    questions = get_test_questions()
    logger.info(f"使用 {len(questions)} 个测试问题")
    
    results = {
        'SentenceTransformer': test_sentence_transformer_load_questions(),
        'BertChinese': test_bert_chinese_load_questions(),
        'BertSimilarity': test_bert_similarity_load_questions()
    }
    
    # 添加聚类检索器测试（如果存在）
    try:
        from questionretrieval import DBSCANClusteringRetriever
        def test_dbscan_clustering_load_questions():
            logger.info("\n=== 测试 DBSCAN 聚类检索器加载问题 ===")
            try:
                retriever = DBSCANClusteringRetriever()
                questions = get_test_questions()
                retriever.load_candidates(questions)
                logger.info(f"✓ 成功加载 {len(questions)} 个问题")
                return True
            except Exception as e:
                logger.error(f"✗ 加载问题失败: {e}")
                return False
        results['DBSCANClustering'] = test_dbscan_clustering_load_questions()
    except ImportError:
        pass
    
    try:
        from questionretrieval import HierarchicalClusteringRetriever
        def test_hierarchical_clustering_load_questions():
            logger.info("\n=== 测试层级聚类检索器加载问题 ===")
            try:
                retriever = HierarchicalClusteringRetriever()
                questions = get_test_questions()
                retriever.load_candidates(questions)
                logger.info(f"✓ 成功加载 {len(questions)} 个问题")
                return True
            except Exception as e:
                logger.error(f"✗ 加载问题失败: {e}")
                return False
        results['HierarchicalClustering'] = test_hierarchical_clustering_load_questions()
    except ImportError:
        pass
    
    logger.info("\n" + "-"*40)
    logger.info("问题加载测试结果")
    logger.info("-"*40)
    
    success_count = 0
    for name, success in results.items():
        status = "✓ 成功" if success else "✗ 失败"
        logger.info(f"{name:20} : {status}")
        if success:
            success_count += 1
    
    logger.info(f"\n总计: {success_count}/{len(results)} 个检索器可以加载问题")
    
    if success_count > 0:
        logger.info("\n💡 提示: 可以加载问题的检索器可以进行推荐测试")
    else:
        logger.warning("\n⚠️  没有检索器可以加载问题，请检查初始化是否成功")
    
    return results

if __name__ == "__main__":
    test_all_retrievers_load_questions()